import functools
import logging
import weakref
import webapp2

from google.appengine.api import datastore
from google.appengine.api import memcache
from google.appengine.ext import db


class Cache(object):
  def __init__(self):
    self._cache = {}
    self._dirty = False

  def get(self, key):
    value = self._cache.get(key, None)
    if value:
      return value

    value = memcache.get(key, namespace=self.__class__.__name__)
    if value:
      self._cache[key] = value
      # don't set dirty bit, since we just got it from memcache
      return value

  def get_multi(self, keys):
    results = [self._cache.get(key, None) for key in keys]
    to_fetch = dict((keys[i], i) for (i, result) in enumerate(results) if not result)
    if not to_fetch:
      return results
    values = memcache.get_multi(to_fetch.iterkeys(), namespace=self.__class__.__name__)
    for key, value in values.iteritems():
      self._cache[key] = value
      results[to_fetch[key]] = value
    return results

  def set(self, key, value):
    if self._cache.get(key, None) == value:
      return
    self._cache[key] = value
    self._dirty = True

  def save_to_memcache(self):
    if not self._dirty:
      return
    memcache.set_multi(self._cache, namespace=self.__class__.__name__)
    self._dirty = False

  def clear(self):
    self._cache.clear()


_model_cache = Cache()
_query_cache = Cache()
_response_cache = Cache()


class Model(db.Model):
  @classmethod
  def from_entity(cls, entity):
    model = super(Model, cls).from_entity(entity)
    if model:
      model.save_to_cache()
    return model

  def save_to_cache(self):
    _model_cache.set(str(self.key()), self)


def _IterWrapper(key, real_iter):
  items = []
  for item in real_iter:
    items.append(item)
    yield item
  _query_cache.set(key, items)


class Query(db.Query):
  def run(self, *args, **kwargs):
    key = str((self._model_class, self._get_query()))
    result = _query_cache.get(key)
    if result:
      for item in result:
        if isinstance(item, Model):
          item.save_to_cache()
      return iter(result)

    real_iter = db.Query.run(self, *args, **kwargs)
    return _IterWrapper(key, real_iter)


def QueryMulti(queries):
  values = _query_cache.get_multi(
      [str((query._model_class, query._get_query()))
       for query in queries])
  for i, value in enumerate(values):
    if not value:
      key = str((queries[i]._model_class, queries[i]._get_query()))
      values[i] = _IterWrapper(key, queries[i].run())
  return values


class RequestHandler(webapp2.RequestHandler):
  def get(self):
    if self.request.get('nocache'):
      return self.get_internal()

    key = str(self.request.params)
    response = _response_cache.get(key)
    if response:
      self.response.out.write(response)
      return

    self.get_internal()

    _response_cache.set(key, self.response.body)


def SaveToMemcache():
  _model_cache.save_to_memcache()
  _query_cache.save_to_memcache()
  _response_cache.save_to_memcache()


def _GetCached(real_func, keys, **kwargs):
  real_keys, multiple = datastore.NormalizeAndTypeCheckKeys(keys)
  if multiple:
    return real_func(keys, **kwargs)

  real_key = str(real_keys[0])

  return (_model_cache.get(real_key) or
          real_func(keys, **kwargs))


_original_db_get = db.get
db.get = functools.partial(_GetCached, db.get)


def DisableCache():
  db.get = _original_db_get


def FlushLocalCache():
  """Flush the cache in just this instance."""
  _model_cache.clear()
  _query_cache.clear()
  _response_cache.clear()


def uses_cache(handler):

  @functools.wraps(handler)
  def SaveCache(self):
    handler(self)
    SaveToMemcache()

  return SaveCache
